package memory

import (
    "strings"
    "sync"
    "time"
    
    "github.com/yangyongzhen/mqtt-broker/internal/protocol/common"
    "github.com/yangyongzhen/mqtt-broker/internal/storage"
)

type MemoryStore struct {
    sessions        map[string]*storage.Session
    retainedMsgs    map[string]*common.Message
    clientMessages  map[string][]*common.Message
    mu              sync.RWMutex
}

func NewMemoryStore() *MemoryStore {
    return &MemoryStore{
        sessions:       make(map[string]*storage.Session),
        retainedMsgs:   make(map[string]*common.Message),
        clientMessages: make(map[string][]*common.Message),
    }
}

func (m *MemoryStore) SaveSession(clientID string, session *storage.Session) error {
    m.mu.Lock()
    defer m.mu.Unlock()
    
    // 深拷贝会话
    sessionCopy := &storage.Session{
        ClientID:      session.ClientID,
        CleanSession:  session.CleanSession,
        Subscriptions: make(map[string]*common.Subscription),
        PendingAcks:   make(map[uint16]*common.Message),
        LastSeen:      session.LastSeen,
    }
    
    for topic, sub := range session.Subscriptions {
        sessionCopy.Subscriptions[topic] = &common.Subscription{
            Topic: sub.Topic,
            QoS:   sub.QoS,
        }
    }
    
    for packetID, msg := range session.PendingAcks {
        sessionCopy.PendingAcks[packetID] = &common.Message{
            Topic:     msg.Topic,
            Payload:   append([]byte(nil), msg.Payload...),
            QoS:       msg.QoS,
            Retain:    msg.Retain,
            PacketID:  msg.PacketID,
            Timestamp: msg.Timestamp,
        }
    }
    
    m.sessions[clientID] = sessionCopy
    return nil
}

func (m *MemoryStore) LoadSession(clientID string) (*storage.Session, error) {
    m.mu.RLock()
    defer m.mu.RUnlock()
    
    session, exists := m.sessions[clientID]
    if !exists {
        return nil, nil
    }
    
    // 深拷贝返回
    sessionCopy := &storage.Session{
        ClientID:      session.ClientID,
        CleanSession:  session.CleanSession,
        Subscriptions: make(map[string]*common.Subscription),
        PendingAcks:   make(map[uint16]*common.Message),
        LastSeen:      session.LastSeen,
    }
    
    for topic, sub := range session.Subscriptions {
        sessionCopy.Subscriptions[topic] = &common.Subscription{
            Topic: sub.Topic,
            QoS:   sub.QoS,
        }
    }
    
    for packetID, msg := range session.PendingAcks {
        sessionCopy.PendingAcks[packetID] = &common.Message{
            Topic:     msg.Topic,
            Payload:   append([]byte(nil), msg.Payload...),
            QoS:       msg.QoS,
            Retain:    msg.Retain,
            PacketID:  msg.PacketID,
            Timestamp: msg.Timestamp,
        }
    }
    
    return sessionCopy, nil
}

func (m *MemoryStore) DeleteSession(clientID string) error {
    m.mu.Lock()
    defer m.mu.Unlock()
    
    delete(m.sessions, clientID)
    delete(m.clientMessages, clientID)
    return nil
}

func (m *MemoryStore) SaveMessage(clientID string, message *common.Message) error {
    m.mu.Lock()
    defer m.mu.Unlock()
    
    // 深拷贝消息
    msgCopy := &common.Message{
        Topic:     message.Topic,
        Payload:   append([]byte(nil), message.Payload...),
        QoS:       message.QoS,
        Retain:    message.Retain,
        PacketID:  message.PacketID,
        Timestamp: message.Timestamp,
    }
    
    m.clientMessages[clientID] = append(m.clientMessages[clientID], msgCopy)
    return nil
}

func (m *MemoryStore) LoadMessages(clientID string) ([]*common.Message, error) {
    m.mu.RLock()
    defer m.mu.RUnlock()
    
    messages, exists := m.clientMessages[clientID]
    if !exists {
        return nil, nil
    }
    
    // 深拷贝返回
    result := make([]*common.Message, len(messages))
    for i, msg := range messages {
        result[i] = &common.Message{
            Topic:     msg.Topic,
            Payload:   append([]byte(nil), msg.Payload...),
            QoS:       msg.QoS,
            Retain:    msg.Retain,
            PacketID:  msg.PacketID,
            Timestamp: msg.Timestamp,
        }
    }
    
    return result, nil
}

func (m *MemoryStore) DeleteMessage(clientID string, packetID uint16) error {
    m.mu.Lock()
    defer m.mu.Unlock()
    
    messages, exists := m.clientMessages[clientID]
    if !exists {
        return nil
    }
    
    // 删除指定包ID的消息
    for i, msg := range messages {
        if msg.PacketID == packetID {
            m.clientMessages[clientID] = append(messages[:i], messages[i+1:]...)
            break
        }
    }
    
    return nil
}

func (m *MemoryStore) SaveRetainedMessage(topic string, message *common.Message) error {
    m.mu.Lock()
    defer m.mu.Unlock()
    
    if message == nil || len(message.Payload) == 0 {
        delete(m.retainedMsgs, topic)
    } else {
        // 深拷贝消息
        msgCopy := &common.Message{
            Topic:     message.Topic,
            Payload:   append([]byte(nil), message.Payload...),
            QoS:       message.QoS,
            Retain:    message.Retain,
            PacketID:  message.PacketID,
            Timestamp: message.Timestamp,
        }
        m.retainedMsgs[topic] = msgCopy
    }
    
    return nil
}

func (m *MemoryStore) LoadRetainedMessage(topic string) (*common.Message, error) {
    m.mu.RLock()
    defer m.mu.RUnlock()
    
    msg, exists := m.retainedMsgs[topic]
    if !exists {
        return nil, nil
    }
    
    // 深拷贝返回
    return &common.Message{
        Topic:     msg.Topic,
        Payload:   append([]byte(nil), msg.Payload...),
        QoS:       msg.QoS,
        Retain:    msg.Retain,
        PacketID:  msg.PacketID,
        Timestamp: msg.Timestamp,
    }, nil
}

func (m *MemoryStore) DeleteRetainedMessage(topic string) error {
    m.mu.Lock()
    defer m.mu.Unlock()
    
    delete(m.retainedMsgs, topic)
    return nil
}

func (m *MemoryStore) SaveSubscription(clientID string, subscription *common.Subscription) error {
    m.mu.Lock()
    defer m.mu.Unlock()
    
    session, exists := m.sessions[clientID]
    if !exists {
        session = &storage.Session{
            ClientID:      clientID,
            Subscriptions: make(map[string]*common.Subscription),
            PendingAcks:   make(map[uint16]*common.Message),
            LastSeen:      time.Now(),
        }
        m.sessions[clientID] = session
    }
    
    session.Subscriptions[subscription.Topic] = &common.Subscription{
        Topic: subscription.Topic,
        QoS:   subscription.QoS,
    }
    
    return nil
}

func (m *MemoryStore) LoadSubscriptions(clientID string) ([]*common.Subscription, error) {
    m.mu.RLock()
    defer m.mu.RUnlock()
    
    session, exists := m.sessions[clientID]
    if !exists {
        return nil, nil
    }
    
    var subscriptions []*common.Subscription
    for _, sub := range session.Subscriptions {
        subscriptions = append(subscriptions, &common.Subscription{
            Topic: sub.Topic,
            QoS:   sub.QoS,
        })
    }
    
    return subscriptions, nil
}

func (m *MemoryStore) DeleteSubscription(clientID string, topic string) error {
    m.mu.Lock()
    defer m.mu.Unlock()
    
    session, exists := m.sessions[clientID]
    if !exists {
        return nil
    }
    
    delete(session.Subscriptions, topic)
    return nil
}

// 扩展方法：获取所有保留消息
func (m *MemoryStore) GetAllRetainedMessages() map[string]*common.Message {
    m.mu.RLock()
    defer m.mu.RUnlock()
    
    result := make(map[string]*common.Message)
    for topic, msg := range m.retainedMsgs {
        result[topic] = &common.Message{
            Topic:     msg.Topic,
            Payload:   append([]byte(nil), msg.Payload...),
            QoS:       msg.QoS,
            Retain:    msg.Retain,
            PacketID:  msg.PacketID,
            Timestamp: msg.Timestamp,
        }
    }
    
    return result
}

// 扩展方法：获取匹配主题过滤器的保留消息
func (m *MemoryStore) GetMatchingRetainedMessages(topicFilter string) []*common.Message {
    m.mu.RLock()
    defer m.mu.RUnlock()
    
    var messages []*common.Message
    
    for topic, msg := range m.retainedMsgs {
        if m.topicMatch(topic, topicFilter) {
            messages = append(messages, &common.Message{
                Topic:     msg.Topic,
                Payload:   append([]byte(nil), msg.Payload...),
                QoS:       msg.QoS,
                Retain:    msg.Retain,
                PacketID:  msg.PacketID,
                Timestamp: msg.Timestamp,
            })
        }
    }
    
    return messages
}

func (m *MemoryStore) topicMatch(topic, filter string) bool {
    if filter == topic {
        return true
    }
    
    filterParts := strings.Split(filter, "/")
    topicParts := strings.Split(topic, "/")
    
    return m.matchParts(topicParts, filterParts, 0, 0)
}

func (m *MemoryStore) matchParts(topic, filter []string, ti, fi int) bool {
    if fi >= len(filter) {
        return ti >= len(topic)
    }
    
    if filter[fi] == "#" {
        return true
    }
    
    if ti >= len(topic) {
        return false
    }
    
    if filter[fi] == "+" || filter[fi] == topic[ti] {
        return m.matchParts(topic, filter, ti+1, fi+1)
    }
    
    return false
}
